package org.ruoyi.interview.config;

import java.io.*;
import java.net.*;
import java.nio.charset.StandardCharsets;
import com.google.gson.*;

public class ModelScopeClient {
    // ModelScope API基础地址 (根据ModelScope API文档)
    private static final String BASE_URL = "https://api-inference.modelscope.cn";
    // 备用API基础地址（使用IP地址）
    private static final String FALLBACK_BASE_URL = "https://47.110.138.66";
    // 你的API密钥
    private final String apiKey;
    // HTTP连接超时时间(毫秒)
    private static final int CONNECT_TIMEOUT = 5000;
    // 读取响应超时时间(毫秒)
    private static final int READ_TIMEOUT = 30000;

    public ModelScopeClient(String apiKey) {
        this.apiKey = apiKey;
    }

    /**
     * 调用文本生成API
     * @param model 模型名称
     * @param prompt 输入提示
     * @return 生成的文本结果
     * @throws IOException 网络请求异常
     */
    public String generateText(String model, String prompt) throws IOException {
        String url = BASE_URL + "/api-inference/v1/models/" + model;

        // 构建请求体
        JsonObject requestBody = new JsonObject();
        JsonObject inputs = new JsonObject();
        inputs.addProperty("text", prompt);
        requestBody.add("inputs", inputs);
        JsonObject parameters = new JsonObject();
        requestBody.add("parameters", parameters);

        String jsonResponse = sendPostRequest(url, requestBody.toString());

        // 解析响应
        JsonElement responseElement = com.google.gson.JsonParser.parseString(jsonResponse);
        if (responseElement.isJsonArray()) {
            JsonArray responseArray = responseElement.getAsJsonArray();
            if (responseArray.size() > 0) {
                JsonObject responseObject = responseArray.get(0).getAsJsonObject();
                if (responseObject.has("generated_text")) {
                    return responseObject.get("generated_text").getAsString();
                }
            }
        } else if (responseElement.isJsonObject()) {
            JsonObject responseObject = responseElement.getAsJsonObject();
            if (responseObject.has("generated_text")) {
                return responseObject.get("generated_text").getAsString();
            }
        }

        throw new IOException("Failed to get valid response from ModelScope API: " + jsonResponse);
    }

    /**
     * 发送POST请求
     * @param url 请求地址
     * @param jsonBody 请求体JSON字符串
     * @return 响应字符串
     * @throws IOException 网络请求异常
     */
    private String sendPostRequest(String url, String jsonBody) throws IOException {
        // 检查是否可以解析域名，如果不能则使用备用IP地址
        String finalUrl = resolveUrlWithFallback(url);

        URL apiUrl = new URL(finalUrl);
        HttpURLConnection connection = (HttpURLConnection) apiUrl.openConnection();

        try {
            // 设置请求方法和头部信息
            connection.setRequestMethod("POST");
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setRequestProperty("Authorization", "Bearer " + apiKey);
            connection.setConnectTimeout(CONNECT_TIMEOUT);
            connection.setReadTimeout(READ_TIMEOUT);

            // 允许输出
            connection.setDoOutput(true);

            // 发送请求体
            try (OutputStream os = connection.getOutputStream()) {
                byte[] input = jsonBody.getBytes(StandardCharsets.UTF_8);
                os.write(input, 0, input.length);
            }

            // 读取响应
            int responseCode = connection.getResponseCode();
            if (responseCode != HttpURLConnection.HTTP_OK) {
                // 读取错误响应体
                StringBuilder errorResponse = new StringBuilder();
                try (BufferedReader br = new BufferedReader(
                        new InputStreamReader(connection.getErrorStream(), StandardCharsets.UTF_8))) {
                    String responseLine;
                    while ((responseLine = br.readLine()) != null) {
                        errorResponse.append(responseLine.trim());
                    }
                } catch (Exception e) {
                    // 忽略错误读取异常
                }
                throw new IOException("HTTP request failed with response code: " + responseCode + 
                                    ", error response: " + errorResponse.toString());
            }

            try (BufferedReader br = new BufferedReader(
                new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8))) {
                StringBuilder response = new StringBuilder();
                String responseLine;
                while ((responseLine = br.readLine()) != null) {
                    response.append(responseLine.trim());
                }
                return response.toString();
            }
        } finally {
            connection.disconnect();
        }
    }

    /**
     * 解析URL，如果DNS解析失败则使用备用IP地址
     *
     * @param originalUrl 原始URL
     * @return 解析后的URL
     * @throws IOException URL处理异常
     */
    private String resolveUrlWithFallback(String originalUrl) throws IOException {
        try {
            // 尝试解析域名
            String host = "api-inference.modelscope.cn";
            InetAddress address = InetAddress.getByName(host);
            System.out.println("域名解析成功: " + host + " -> " + address.getHostAddress());
            return originalUrl;
        } catch (UnknownHostException e) {
            System.out.println("域名解析失败: " + e.getMessage() + ", 尝试使用备用IP地址");
            // 使用备用API地址替换原始URL中的域名部分
            String fallbackUrl = originalUrl.replace("https://api-inference.modelscope.cn", 
                                                   "https://47.110.138.66");
            System.out.println("使用备用API地址: " + fallbackUrl);
            return fallbackUrl;
        }
    }

    // 使用示例
    public static void main(String[] args) {
        // 替换为你的API密钥
        String apiKey = "ms-b4cce3b6-9a9d-4fb1-9713-ccb77ba72934";
        ModelScopeClient client = new ModelScopeClient(apiKey);

        try {
            // 调用文本生成API，例如使用"baichuan2-7b-chat"模型
            String result = client.generateText("baichuan2-7b-chat", "请介绍一下人工智能的发展历程");
            System.out.println("生成结果: " + result);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}